#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 功能：ROS2 双目深度节点，利用 SGBM + WLS 生成稠密视差图，
#       并用 ORB 特征匹配计算部分关键点距离，实时可视化。

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import cv2
import numpy as np
import yaml


class StereoDepth(Node):
    def __init__(self):
        super().__init__('stereo_depth')
        self.bridge = CvBridge()

        # 订阅左右原始图像话题
        self.left_sub  = self.create_subscription(
            Image, '/left/image_raw',  self.left_cb,  10)
        self.right_sub = self.create_subscription(
            Image, '/right/image_raw', self.right_cb, 10)

        # 缓存最新左右图
        self.left_img = self.right_img = None

        # 加载标定参数并生成矫正映射
        self.load_calib()
        # 创建 SGBM 立体匹配器 + WLS 滤波器
        self.setup_stereo_matcher()

    # ------------------------------------------------------------------
    # ① 读取 YAML 标定文件，计算立体矫正映射
    # ------------------------------------------------------------------
    def load_calib(self):
        with open('install/stereo_depth_py/share/stereo_depth_py/config/stereo.yaml') as f:
            cfg = yaml.safe_load(f)

        # 相机内参、畸变、外参
        self.K1 = np.array(cfg['K1'])  # 左相机内参
        self.D1 = np.array(cfg['D1'])  # 左相机畸变
        self.K2 = np.array(cfg['K2'])  # 右相机内参
        self.D2 = np.array(cfg['D2'])  # 右相机畸变
        self.R  = np.array(cfg['R'])   # 旋转矩阵
        self.T  = np.array(cfg['T'])   # 平移向量 (基线)

        # 基线长度 fx 用于三角测距
        self.baseline = abs(self.T[0])
        self.fx = self.K1[0, 0]

        h, w = 480, 640
        # 计算立体矫正参数
        self.R1, self.R2, self.P1, self.P2, self.Q, *_ = cv2.stereoRectify(
            self.K1, self.D1, self.K2, self.D2, (w, h), self.R, self.T, alpha=0)

        # 生成矫正映射表（remap 用）
        self.map1x, self.map1y = cv2.initUndistortRectifyMap(
            self.K1, self.D1, self.R1, self.P1, (w, h), cv2.CV_32FC1)
        self.map2x, self.map2y = cv2.initUndistortRectifyMap(
            self.K2, self.D2, self.R2, self.P2, (w, h), cv2.CV_32FC1)

    # ------------------------------------------------------------------
    # ② 创建立体匹配器：SGBM + WLS 滤波
    # ------------------------------------------------------------------
    def setup_stereo_matcher(self):
        self.sgbm = cv2.StereoSGBM_create(
            minDisparity=0,          # 最小视差
            numDisparities=128,      # 视差搜索范围，必须是16整数倍
            blockSize=5,             # 匹配块大小
            P1=8 * 3 * 5 ** 2,       # 平滑项系数
            P2=32 * 3 * 5 ** 2,
            disp12MaxDiff=1,
            uniquenessRatio=15,
            speckleWindowSize=100,
            speckleRange=1,
            mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY
        )

        # WLS 滤波器，用于边缘保持平滑
        self.wls = cv2.ximgproc.createDisparityWLSFilter(self.sgbm)
        self.wls.setLambda(8000)
        self.wls.setSigmaColor(1.5)

    # ------------------------------------------------------------------
    # ③ 左右图回调：右图到达时触发一次完整处理
    # ------------------------------------------------------------------
    def left_cb(self, msg):
        self.left_img = self.bridge.imgmsg_to_cv2(msg, 'bgr8')

    def right_cb(self, msg):
        self.right_img = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
        if self.left_img is None:
            return
        self.process()

    # ------------------------------------------------------------------
    # ④ 主要处理流程：矫正 → 视差计算 → 特征点测距 → 可视化
    # ------------------------------------------------------------------
    def process(self):
        # 1. 立体矫正
        l = cv2.remap(self.left_img,  self.map1x, self.map1y, cv2.INTER_LINEAR)
        r = cv2.remap(self.right_img, self.map2x, self.map2y, cv2.INTER_LINEAR)

        # 2. 计算视差
        disp = self.sgbm.compute(l, r).astype(np.float32) / 16.0  # 转为像素视差
        disp = self.wls.filter(disp, l)                            # 边缘保持滤波

        # 3. ORB 特征匹配，用于稀疏深度验证
        orb = cv2.ORB_create(200)
        gray_l = cv2.cvtColor(l, cv2.COLOR_BGR2GRAY)
        gray_r = cv2.cvtColor(r, cv2.COLOR_BGR2GRAY)
        kp1, des1 = orb.detectAndCompute(gray_l, None)
        kp2, des2 = orb.detectAndCompute(gray_r, None)

        bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        matches = bf.match(des1, des2)

        # 计算前 20 个匹配点的距离
        for m in matches[:20]:
            pt1 = kp1[m.queryIdx].pt  # 左图坐标
            pt2 = kp2[m.trainIdx].pt  # 右图坐标
            d = abs(pt1[0] - pt2[0])  # 视差（像素）
            if d > 0:
                # 三角测距：Z = (b * f) / d
                Z = (self.baseline * self.fx) / d
                # 画圆+文字
                cv2.circle(l, tuple(map(int, pt1)), 4, (0, 255, 0), -1)
                cv2.putText(l, f'{Z:.2f}m',
                            (int(pt1[0]), int(pt1[1]) - 5),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.5, (0, 0, 255), 1)

        # 4. 实时显示
        cv2.imshow('disparity', disp / 128.0)  # 归一化到 0~1
        cv2.imshow('left', l)
        cv2.waitKey(1)


# ------------------------------------------------------------------
# ⑤ ROS2 节点入口
# ------------------------------------------------------------------
def main():
    rclpy.init()
    rclpy.spin(StereoDepth())
    cv2.destroyAllWindows()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
